import numpy as np
from scipy import signal
import matplotlib.pyplot as plt

# ======================
# 参数配置区（按需修改）
# ======================
# 试验参数
sampling_rate = 5000  # 采样率 (Hz)
duration = 1.0  # 信号时长 (秒)
n_points = int(sampling_rate * duration)  # 总采样点数
n_measures = 6  # 测点数量（K3、K8等）

# 滤波器参数
filter_method2_cutoff = 300  # 方法2低通滤波截止频率 (Hz)
filter_method3_band = [2, 500]  # 方法3带通滤波范围 (Hz)

# 时间窗参数
method1_window = 0.1  # 方法1分析时间窗长度 (秒)
method3_delta_t = 1.0  # 方法3积分时间间隔 (秒)

# ======================
# 模拟数据生成（替换为实际数据）
# ======================
t = np.linspace(0, duration, n_points)
# 生成示例压力信号（实际需替换为试验数据）
pressure_signals = [
    100 * np.sin(2 * np.pi * 50 * t) + 10 * np.random.randn(n_points) + 500  # 测点1
    for _ in range(n_measures)
]


# ======================
# 数据处理方法实现
# ======================

def method1_peak_to_peak(signal, window_time):
    """
    方法1：计算脉动压力峰峰值
    :param signal: 输入时域信号 (1D numpy array)
    :param window_time: 分析时间窗长度 (秒)
    :return: 峰峰值 (单位同输入信号)
    """
    window_samples = int(window_time * sampling_rate)
    peak = np.max(signal[-window_samples:])  # 取最后window_time秒数据
    valley = np.min(signal[-window_samples:])
    return peak - valley


def method2_rms_fluctuation(signal):
    """
    方法2：计算脉动量（带300Hz低通滤波）
    :param signal: 输入时域信号 (1D numpy array)
    :return: 脉动量RMS值
    """
    # 设计6阶巴特沃斯低通滤波器
    nyquist = 0.5 * sampling_rate
    normalized_cutoff = filter_method2_cutoff / nyquist
    b, a = signal.butter(6, normalized_cutoff, btype='low')

    # 滤波处理（零相位滤波）
    filtered_signal = signal.filtfilt(b, a, signal)

    # 计算脉动量
    mean_val = np.mean(filtered_signal)
    rms = np.sqrt(np.mean((filtered_signal - mean_val) ** 2))
    return rms

def method3_turbulence_intensity(signals):
    """
    方法3：计算面平均紊流度（带2-500Hz带通滤波）
    :param signals: 多测点信号列表 (list of 1D numpy arrays)
    :return: 面平均紊流度
    """
    # 设计4阶巴特沃斯带通滤波器
    nyquist = 0.5 * sampling_rate
    low = filter_method3_band[0] / nyquist
    high = filter_method3_band[1] / nyquist
    b, a = signal.butter(4, [low, high], btype='band')

    # 对各测点信号进行滤波
    filtered_signals = [signal.filtfilt(b, a, s) for s in signals]

    # 计算各测点脉动量RMS
    rms_values = [
        np.sqrt(np.mean((s - np.mean(s)) ** 2))
        for s in filtered_signals
    ]

    # 面平均计算
    integral_time = method3_delta_t * sampling_rate  # 转换为样本数
    avg_rms = np.sqrt(np.mean(np.square(rms_values)))
    return avg_rms


# ======================
# 执行计算并输出结果
# ======================
if __name__ == "__main__":
    # 方法1计算（取最后一个测点示例）
    p1_pp = method1_peak_to_peak(pressure_signals[0], method1_window)
    print(f"方法1 - K3截面脉动峰峰值: {p1_pp:.2f} kPa")

    # 方法2计算（取第一个测点示例）
    p2_rms = method2_rms_fluctuation(pressure_signals[0])
    print(f"方法2 - K3截面脉动量RMS: {p2_rms:.2f} kPa")

    # 方法3计算（全部测点）
    turbulence = method3_turbulence_intensity(pressure_signals)
    print(f"方法3 - AIP截面紊流度: {turbulence:.3f}")

# ======================
# 结果可视化（可选）
# ======================
plt.figure(figsize=(12, 8))

# 绘制原始信号与处理后信号
plt.subplot(3, 1, 1)
plt.plot(t, pressure_signals[0], label='原始信号')
plt.ylabel('Pressure (kPa)')
plt.legend()

plt.subplot(3, 1, 2)
plt.plot(t, method2_rms_fluctuation(pressure_signals[0]),
         label='方法2滤波后信号', color='orange')
plt.ylabel('Pressure (kPa)')
plt.legend()

plt.subplot(3, 1, 3)
plt.plot(t, pressure_signals[0], label='原始信号')
plt.plot(t, method3_turbulence_intensity([pressure_signals[0]]) * np.ones_like(t),
         label='方法3紊流度', color='green')
plt.xlabel('Time (s)')
plt.ylabel('Pressure (kPa)')
plt.legend()

plt.tight_layout()
plt.show()